from langgraph.graph import StateGraph
from typing_extensions import TypedDict
from langgraph.graph import START,END
'''''状态图定义——————StateGraph'''
#定义输入状态的格式
class InputState(TypedDict):
  question: str

#定义输出状态的格式
class OutputState(TypedDict):
  answer: str

#将InputState和OutputState两个TypedDict合并为一个字典类型
class OverallState(InputState, OutputState):
  pass

#创建builder对象，完成状态图框架构建
builder = StateGraph(OverallState,input=InputState,output=OutputState)

'''图节点定义————————Nodes'''

#定义模拟智能体节点
def agent_node(state: InputState):
  print('我是一个智能体节点')
  return

#定义模拟行动节点
def action_node(state: InputState):
  print('我是一个行动节点')
  return {'answer':'我已经执行完毕了'}

#添加节点到状态图中
builder.add_node('智能体',agent_node)
builder.add_node('执行者',action_node)

'''图边定义————————Edges'''
#添加起始节点
builder.add_edge(START, '智能体')
#添加智能体节点到行动节点的边
builder.add_edge('智能体', '执行者')
#添加行动节点到结束节点的边
builder.add_edge('执行者', END)

'''状态图构建完成，激活状态图'''
graph = builder.compile()
# graph.invoke({'question':'请你帮我写一段代码'})
output = graph.invoke({'question':'请你帮我写一段代码'})
# print(output)
